import os
import re
import ast
import csv


def array_intersection_length(arr1, arr2, v):
    # 返回两个集合的长度，交集长度，以及lambda
    set1 = set(arr1)
    set2 = set(arr2)
    intersection = set1.intersection(set2)

    k1 = len(arr1)
    k2 = len(arr2)
    e = len(intersection)
    lambd = int((k1*k2-e)/(v-1))

    return k1, k2, e, lambd


def extract_integers_from_filename(filename):
    # 获取group和index
    match = re.match(r'(\d+)-(\d+)\.txt', filename)
    if match:
        int1 = int(match.group(1))
        int2 = int(match.group(2))
        return int1, int2
    else:
        raise ValueError


def read_text_file(directory_path, filename):
    # 读取txt文件，返回pairs
    file_path = os.path.join(directory_path, filename)
    pairs = []
    sum = 0
    with open(file_path, 'r') as file:
        # skip the first row
        next(file)

        # read in the amout of the pairs
        second_line = next(file)
        amout = int(second_line.split(":")[1].strip())

        # skip line 3 ~ 5
        for _ in range(3):
            next(file)

        # read lines and concatenate pairs
        pair_lines = []
        
        for line in file:
            line = line.strip()
            if line == '];' or line == '':
                continue

            if line.startswith("[ [ "):
                pair_str = ''.join(pair_lines)
                if pair_str.endswith(','):
                    pair_str = pair_str[:-1]
                if pair_str != '':
                    pair = ast.literal_eval(pair_str)
                    pairs.append(pair)
                    sum = sum + 1
                pair_lines = []
                pair_lines.append(line)
            elif line.endswith("] ],") or line.endswith(" ] ]"):
                pair_lines.append(line)
                pair_str = ''.join(pair_lines)
                if pair_str.endswith(','):
                    pair_str = pair_str[:-1]
                pair = ast.literal_eval(pair_str)
                pairs.append(pair)
                sum = sum + 1
                pair_lines = []
            else:
                pair_lines.append(line)

        pair_str = ''.join(pair_lines)
        if pair_str.endswith(','):
            pair_str = pair_str[:-1]
        if pair_str != '':
            pair = ast.literal_eval(pair_str)
            pairs.append(pair)
            sum = sum + 1  

        # check if the amout is correct
        if (sum != amout):
            print("Error, the amout is incorrect!!!")
        return pairs


def write_to_csv(order, index, pairs, path_to_csv):
    # 写到文件中
    with open(path_to_csv, 'w', newline='') as file:
        csv_writer = csv.writer(file)
        csv_writer.writerow(
            ['order', 'index', 'v', 'k1', 'k2', 'e', 'lambda', 'n', 'dsp'])
        for dsp in pairs:
            k1, k2, e, lambd = array_intersection_length(dsp[0], dsp[1], order)
            n = e - lambd
            csv_writer.writerow([order, index, order, k1, k2, e, lambd, n, dsp])


def process_txt(directory_path):
    txt_files = [file for file in os.listdir(directory_path) if file.endswith(".txt")]
    for file in txt_files:
        pairs = read_text_file(directory_path, file)
        if len(pairs) == 0 :
            continue
        order, index = extract_integers_from_filename(file)
        write_to_csv(order, index, pairs,f"{order}-{index}.csv")

process_txt("..")
